import os

import pandas as pd

from app_config import get_engine_ts, get_pro

"""
csi all exchange_amount > 300 yi
exclude hs_300 
exclude zz_500
exclude sz_50
"""
engine = get_engine_ts()


def calc(date_sz50: str, strStartDate: str, strEndDate: str):
    df_csi_all = get_pro().index_weight(index_code='000985.CSI', start_date=date_sz50, end_date=date_sz50)
    list_csi_all_code = df_csi_all['con_code'].tolist()

    query = f"""
       SELECT ts_code,trade_date, amount FROM `daily`
       WHERE ts_code IN ({','.join(f"'{code}'" for code in list_csi_all_code)})
       AND trade_date >= '{strStartDate}'
       AND trade_date <= '{strEndDate}'
       """

    # 执行查询并将结果转换为DataFrame
    df_daily_csi_all = pd.read_sql_query(query, engine)
    df_daily_csi_all['amount'] = df_daily_csi_all['amount'] / 1000
    grouped_df = df_daily_csi_all.groupby('ts_code')['amount'].mean().reset_index()
    grouped_df.columns = ['ts_code', 'avg_amount']
    merge = pd.merge(df_csi_all, grouped_df, left_on='con_code', right_on='ts_code', how='left')
    # merge.to_excel('./csi_all.xlsx', index=False)

    # 构建查询语句
    query = f"""
          SELECT ts_code,trade_date, total_mv FROM `daily_basic`
          WHERE ts_code IN ({','.join(f"'{code}'" for code in list_csi_all_code)})
          AND trade_date >= '{strStartDate}'
          AND trade_date <= '{strEndDate}'
          """
    # 执行查询并将结果转换为DataFrame
    result_df = pd.read_sql_query(query, engine)
    result_df['total_mv'] = result_df['total_mv'] / 10000
    grouped_df_total_mv = result_df.groupby('ts_code')['total_mv'].mean().reset_index()
    grouped_df_total_mv.columns = ['ts_code', 'avg_total_mv']
    pd_merge = pd.merge(merge, grouped_df_total_mv, left_on='con_code', right_on='ts_code', how='left')
    pd_merge.to_excel('./csi_all1.xlsx', index=False)


def calc1(date_sz50: str):
    df_csi_all = pd.read_excel('./csi_all1.xlsx')
    df_stock_basic = df_csi_all[((df_csi_all['avg_total_mv'] > 200.0) & (True))]
    # df_stock_basic.to_excel('./stock_basic.xlsx', index=False)
    index_basics = ['000016.SH', '000300.SH', '399905.SZ', '000510.CSI', '399006.SZ', '399673.SZ', '000688.SH', '000852.SH']

    df_sse_50 = get_pro().index_weight(index_code='000016.SH', start_date=date_sz50, end_date=date_sz50)
    df_stock_basic = df_stock_basic[~df_stock_basic['con_code'].isin(df_sse_50['con_code'])]

    df_csi_300 = get_pro().index_weight(index_code='000300.SH', start_date=date_sz50, end_date=date_sz50)
    df_stock_basic = df_stock_basic[~df_stock_basic['con_code'].isin(df_csi_300['con_code'])]

    df_zz_500 = get_pro().index_weight(index_code='399905.SZ', start_date=date_sz50, end_date=date_sz50)
    df_stock_basic = df_stock_basic[~df_stock_basic['con_code'].isin(df_zz_500['con_code'])]

    df_zz_a500 = get_pro().index_weight(index_code='000510.CSI', start_date=date_sz50, end_date=date_sz50)
    df_stock_basic = df_stock_basic[~df_stock_basic['con_code'].isin(df_zz_a500['con_code'])]

    # 399006.SZ 创业板指
    df_399006_SZ = get_pro().index_weight(index_code='399006.SZ', start_date=date_sz50, end_date=date_sz50)
    df_stock_basic = df_stock_basic[~df_stock_basic['con_code'].isin(df_399006_SZ['con_code'])]
    # 399673.SZ 创业板50
    df_399673_SZ = get_pro().index_weight(index_code='399673.SZ', start_date=date_sz50, end_date=date_sz50)
    df_stock_basic = df_stock_basic[~df_stock_basic['con_code'].isin(df_399673_SZ['con_code'])]
    # 000688.SH 科创50
    df_000688_SH = get_pro().index_weight(index_code='000688.SH', start_date=date_sz50, end_date=date_sz50)
    df_stock_basic = df_stock_basic[~df_stock_basic['con_code'].isin(df_000688_SH['con_code'])]
    # 000852.SH 中证1000
    df_000852_SH = get_pro().index_weight(index_code='000852.SH', start_date=date_sz50, end_date=date_sz50)
    df_stock_basic = df_stock_basic[~df_stock_basic['con_code'].isin(df_000852_SH['con_code'])]

    query = f"""
          SELECT sb.name,sb.ts_code FROM `stock_basic` sb
          """

    # 执行查询并将结果转换为DataFrame
    df_daily_csi_all = pd.read_sql_query(query, engine)

    df_stock_basic.to_excel('./stock_basic.xlsx', index=False)

if __name__ == '__main__':
    calc1(date_sz50="20250430")
